import torch as pt
import numpy as np
from python_ai.common.xcommon import *


class ConvBnRelu(pt.nn.Module):

    def __init__(self, in_ch, out_ch, ksize=3, stride=1, padding=1, bias=False, **kwargs):
        super().__init__(**kwargs)
        self.conv = pt.nn.Conv2d(in_ch, out_ch, ksize, stride, padding, bias=bias)
        self.bn = pt.nn.BatchNorm2d(out_ch)
        self.relu = pt.nn.LeakyReLU(inplace=True, negative_slope=0.1)

    def forward(self, inputs):
        x = self.conv(inputs)
        x = self.bn(x)
        x = self.relu(x)
        return x


cfg = [
    32, 'm',
    64, 'm',
    128, 64, 128, 'm',
    256, 128, 256, 'm',
    512, 256, 512, 256, 512, 'm',
    1024, 512, 1024, 512, 1024,
]


class MyDarknet19(pt.nn.Module):

    def __init__(self, WEIGHT_PATH, in_ch, n_cls, is_pretrained, **kwargs):
        super().__init__(**kwargs)
        layers = []
        ksize = [1, 3]
        ksize_idx = True
        padding = [0, 1]
        padding_idx = True
        for cfg_v in cfg:
            if cfg_v == 'm':
                layers.append(pt.nn.MaxPool2d(2, 2, 0))
            else:
                layers.append(ConvBnRelu(in_ch, cfg_v, ksize[ksize_idx], 1, padding[padding_idx]))
                in_ch = cfg_v
            ksize_idx = not ksize_idx
            padding_idx = not padding_idx
        self.features = pt.nn.Sequential(*layers)
        self.classifier = pt.nn.Sequential(
            pt.nn.Conv2d(cfg_v, n_cls, 1, 1, 0),
            pt.nn.AdaptiveAvgPool2d((1, 1)),
            pt.nn.LogSoftmax(dim=1)
        )

        if is_pretrained:
            dict = {}
            keys = self.state_dict().keys()
            values = pt.load(WEIGHT_PATH).values()
            for k, v in zip(keys, values):
                dict[k] = v
            self.load_state_dict(dict)

    def forward(self, inputs):
        x = self.features(inputs)
        # print('x after features:', x.size())
        x = self.classifier(x)
        # print('x after classifier:', x.size())
        x = pt.squeeze(x, dim=3)
        x = pt.squeeze(x, dim=2)
        return x


criterion = pt.nn.NLLLoss()

if __name__ == '__main__':
    from torch.utils.data import TensorDataset, DataLoader
    import python_ai.common.routine.pt_compile_fit_evaluate as routine
    from python_ai.common.routine.pt_get_device import get_device
    from python_ai.common.read_data.cat_dog_teachers import load_data
    import matplotlib.pyplot as plt

    WEIGHT_PATH = '../../../../../large_data/model/darknet19/darknet19-deepBakSu-e1b3ec1e.pth'

    device = get_device()

    sep('Test')
    net = MyDarknet19(WEIGHT_PATH, in_ch=3, n_cls=1000, is_pretrained=True).to(device)
    print(net)

    x = pt.zeros((2, 3, 224, 224)).to(device)
    out = net(x)
    print(out.size())

    sep('Cat and dog')

    BATCH_SIZE = 8
    N_EPOCHS = 8
    ALPHA = 1e-5

    (x_train, y_train), (x_val, y_val), (x_test, y_test) = load_data(r'../../../../../large_data/DL1/_many_files/zoo',
                                                                     1)

    net = MyDarknet19(WEIGHT_PATH, in_ch=3, n_cls=2, is_pretrained=False).to(device)

    x_train = pt.Tensor(x_train)
    x_test = pt.Tensor(x_test)
    x_val = pt.Tensor(x_val)
    y_train = pt.Tensor(y_train)
    y_test = pt.Tensor(y_test)
    y_val = pt.Tensor(y_val)

    ds_train = TensorDataset(x_train, y_train)
    ds_test = TensorDataset(x_test, y_test)
    ds_val = TensorDataset(x_val, y_val)

    dl_train = DataLoader(ds_train, BATCH_SIZE, drop_last=True)
    dl_test = DataLoader(ds_test, BATCH_SIZE, drop_last=False)
    dl_val = DataLoader(ds_val, BATCH_SIZE, drop_last=False)

    model = net

    model_dict = routine.compile(device, model, criterion, pt.optim.Adam, ALPHA, routine.accuracy)
    loss_his, acc_his, loss_his_val, acc_his_val = routine.fit(model_dict, dl_train, dl_val, N_EPOCHS)

    routine.evaluate(model_dict, dl_test)

    plt.figure(figsize=[12, 6])
    plt_dict = dict(spr=1, spc=2, spn=0)

    routine.plot_loss_and_metric(plt_dict, loss_his, loss_his_val, acc_his, acc_his_val)

    plt.show()
